import os
import subprocess
import winreg
from .base import Tool

def execute(params: dict) -> str:
    """执行system_paths工具"""
    try:
        path_type = params.get("path_type", "").lower()
        if not path_type:
            return "错误：未提供路径类型"
            
        path = get_system_path(path_type)
        if not path:
            return f"错误：无法获取{path_type}路径"
            
        return path
        
    except Exception as e:
        return f"获取系统路径失败: {str(e)}"

TOOL = Tool(
    name="system_paths",
    description="获取Windows系统常用目录路径",
    parameters={
        "path_type": {
            "type": "string",
            "description": "路径类型(desktop/downloads/documents/programs/appdata等)",
            "required": True
        }
    },
    example='''<system_paths>
<path_type>desktop</path_type>
</system_paths>''',
    execute_func=execute
)

SHELL_FOLDERS = {
    "desktop": "Desktop",
    "documents": "Personal",
    "downloads": "{374DE290-123F-4565-9164-39C4925E467B}",
    "music": "My Music",
    "pictures": "My Pictures",
    "videos": "My Video",
    "appdata": "AppData",
    "programs": "Programs",
    "startup": "Startup",
    "recent": "Recent"
}

def get_shell_folder(folder_name: str) -> str:
    """从注册表获取系统文件夹路径"""
    try:
        with winreg.OpenKey(
            winreg.HKEY_CURRENT_USER,
            r"Software\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders"
        ) as key:
            path = winreg.QueryValueEx(key, SHELL_FOLDERS.get(folder_name, folder_name))[0]
            if os.path.exists(path):
                return path
    except:
        pass
    return None

def get_env_path(env_var: str) -> str:
    """从环境变量获取路径"""
    try:
        if os.name == 'nt':  # Windows
            # 使用PowerShell获取环境变量值
            cmd = f'[Environment]::GetFolderPath("{env_var}")'
            result = subprocess.check_output(['powershell', '-Command', cmd], text=True).strip()
            if result and os.path.exists(result):
                return result
    except:
        pass
    return None

def get_system_path(path_type: str) -> str:
    """获取系统路径"""
    # 首先尝试从环境变量获取
    env_map = {
        "desktop": "Desktop",
        "documents": "MyDocuments",
        "music": "MyMusic",
        "pictures": "MyPictures",
        "videos": "MyVideos",
        "programs": "ProgramFiles",
        "appdata": "ApplicationData"
    }
    
    # 尝试从环境变量获取路径
    if path_type in env_map:
        path = get_env_path(env_map[path_type])
        if path:
            return path
    
    # 如果环境变量获取失败，尝试从注册表获取
    path = get_shell_folder(path_type)
    if path:
        return path
        
    # 如果以上都失败，尝试常见默认路径
    user_profile = os.environ.get('USERPROFILE')
    if user_profile:
        default_paths = {
            "desktop": os.path.join(user_profile, "Desktop"),
            "downloads": os.path.join(user_profile, "Downloads"),
            "documents": os.path.join(user_profile, "Documents"),
            "pictures": os.path.join(user_profile, "Pictures"),
            "music": os.path.join(user_profile, "Music"),
            "videos": os.path.join(user_profile, "Videos"),
        }
        
        if path_type in default_paths:
            path = default_paths[path_type]
            if os.path.exists(path):
                return path
    
    return None